/*
 * Licensed to the Apache Software Foundation (ASF) under one or more
 * contributor license agreements.  See the NOTICE file distributed with
 * this work for additional information regarding copyright ownership.
 * The ASF licenses this file to You under the Apache License, Version 2.0
 * (the "License"); you may not use this file except in compliance with
 * the License.  You may obtain a copy of the License at
 *
 *     http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 */
package org.apache.rocketmq.client.consumer.store;

import java.util.HashMap;
import java.util.HashSet;
import java.util.Map;
import java.util.Set;
import java.util.concurrent.ConcurrentHashMap;
import java.util.concurrent.ConcurrentMap;
import java.util.concurrent.atomic.AtomicLong;
import org.apache.rocketmq.client.exception.MQBrokerException;
import org.apache.rocketmq.client.exception.MQClientException;
import org.apache.rocketmq.client.exception.OffsetNotFoundException;
import org.apache.rocketmq.client.impl.FindBrokerResult;
import org.apache.rocketmq.client.impl.factory.MQClientInstance;
import org.apache.rocketmq.common.MixAll;
import org.apache.rocketmq.common.UtilAll;
import org.apache.rocketmq.common.message.MessageQueue;
import org.apache.rocketmq.remoting.exception.RemotingException;
import org.apache.rocketmq.remoting.protocol.header.QueryConsumerOffsetRequestHeader;
import org.apache.rocketmq.remoting.protocol.header.UpdateConsumerOffsetRequestHeader;
import org.apache.rocketmq.logging.org.slf4j.Logger;
import org.apache.rocketmq.logging.org.slf4j.LoggerFactory;

/**
 * Remote storage implementation
 */
public class RemoteBrokerOffsetStore implements OffsetStore {
    private final static Logger log = LoggerFactory.getLogger(RemoteBrokerOffsetStore.class);
    private final MQClientInstance mQClientFactory;
    private final String groupName;
    private ConcurrentMap<MessageQueue, AtomicLong> offsetTable =
        new ConcurrentHashMap<>();

    public RemoteBrokerOffsetStore(MQClientInstance mQClientFactory, String groupName) {
        this.mQClientFactory = mQClientFactory;
        this.groupName = groupName;
    }

    @Override
    public void load() {
    }

    /**
     * 更新偏移量
     * @param mq 指定的MQ
     * @param offset 指定的偏移量，需要保存在broker端
     * @param increaseOnly 是否是只增加的
     */
    @Override
    public void updateOffset(MessageQueue mq, long offset, boolean increaseOnly) {
        if (mq != null) {
            // 从本地偏移量表中根据MQ获取对应的偏移量的原子对象
            AtomicLong offsetOld = this.offsetTable.get(mq);
            if (null == offsetOld) {
                // 如果没有找到，则创建新的原子对象并赋值当前的偏移量，放到偏移量表中
                offsetOld = this.offsetTable.putIfAbsent(mq, new AtomicLong(offset));
            }

            if (null != offsetOld) {
                if (increaseOnly) {
                    // 如果是只增加的，则仅在新值大于现有值的情况下才更新本地缓存
                    MixAll.compareAndIncreaseOnly(offsetOld, offset);
                } else {
                    // 如果不要求必须只增加偏移量的值，则直接设置到本地缓存
                    offsetOld.set(offset);
                }
            }
        }
    }

    /**
     * 读取偏移量：
     * 1. 如果是首先从内存读取，然后从偏移量存储读取，则从本地缓存的偏移量表中读取，如果没有读取到，且读取类型是从内存读取，则直接返回-1；
     * 2. 如果是从远程偏移量读取，则首先从远程读取，读取到之后，更新本地偏移量表的值，返回读取到的值；如果远程没有，则返回-1，其他异常返回-2；
     * 3. 如果读取类型设置其他值，则直接返回-3。
     * @param mq 指定的MQ
     * @param type 从哪里读取
     * @return
     */
    @Override
    public long readOffset(final MessageQueue mq, final ReadOffsetType type) {
        if (mq != null) {
            switch (type) {
                case MEMORY_FIRST_THEN_STORE: // 首先读取内存的，然后从偏移量存储读取
                case READ_FROM_MEMORY: { // 从内存读取
                    // 从本地偏移量表读取
                    AtomicLong offset = this.offsetTable.get(mq);
                    if (offset != null) {
                        // 读取后返回
                        return offset.get();
                    } else if (ReadOffsetType.READ_FROM_MEMORY == type) {
                        // 如果没有读取到，且读取类型是仅从内存读取，则返回-1
                        return -1;
                    }
                }
                case READ_FROM_STORE: { // 从远程偏移量存储读取
                    try {
                        // 远程读取
                        long brokerOffset = this.fetchConsumeOffsetFromBroker(mq);
                        // 更新本地偏移量缓存信息
                        this.updateOffset(mq, brokerOffset, false);
                        // 返回从远程读取到的偏移量
                        return brokerOffset;
                    }
                    // No offset in broker
                    // 如果没有从远程偏移量存储读取到，则返回-1
                    catch (OffsetNotFoundException e) {
                        return -1;
                    }
                    // 其他异常返回-2
                    catch (Exception e) {
                        log.warn("fetchConsumeOffsetFromBroker exception, " + mq, e);
                        return -2;
                    }
                }
                default:
                    break;
            }
        }
        // 如果不是上述分支，则返回-3
        return -3;
    }

    /**
     * 对偏移量执行持久化。
     * 参数表示当前消费者现在所有涉及到的MQ，本地缓存偏移量表中如果有其他的条目，则删除。
     * @param mqs 需要持久化偏移量的MQ集合
     */
    @Override
    public void persistAll(Set<MessageQueue> mqs) {
        // 如果没有指定，则返回
        if (null == mqs || mqs.isEmpty())
            return;

        // 未使用的MQ集合
        final HashSet<MessageQueue> unusedMQ = new HashSet<>();

        // 遍历本地缓存的偏移量表
        for (Map.Entry<MessageQueue, AtomicLong> entry : this.offsetTable.entrySet()) {
            // 获取MQ
            MessageQueue mq = entry.getKey();
            // 获取偏移量
            AtomicLong offset = entry.getValue();
            if (offset != null) {
                if (mqs.contains(mq)) {
                    try {
                        // 将该MQ的偏移量值持久化到远程broker
                        this.updateConsumeOffsetToBroker(mq, offset.get());
                        log.info("[persistAll] Group: {} ClientId: {} updateConsumeOffsetToBroker {} {}",
                            this.groupName,
                            this.mQClientFactory.getClientId(),
                            mq,
                            offset.get());
                    } catch (Exception e) {
                        log.error("updateConsumeOffsetToBroker exception, " + mq.toString(), e);
                    }
                } else {
                    // 如果MQ不是指定的，则缓存到未使用集合中
                    unusedMQ.add(mq);
                }
            }
        }

        // 如果未使用MQ集合中有MQ
        if (!unusedMQ.isEmpty()) {
            // 遍历未使用MQ集合
            for (MessageQueue mq : unusedMQ) {
                // 从本地缓存的偏移量表中移除
                this.offsetTable.remove(mq);
                log.info("remove unused mq, {}, {}", mq, this.groupName);
            }
        }
    }

    /**
     * 对指定的MQ偏移量持久化到远程broker
     * @param mq 需要持久化偏移量的MQ
     */
    @Override
    public void persist(MessageQueue mq) {
        // 获取该MQ的偏移量
        AtomicLong offset = this.offsetTable.get(mq);
        if (offset != null) {
            try {
                // 将指定MQ的偏移量持久化到远程broker
                this.updateConsumeOffsetToBroker(mq, offset.get());
                log.info("[persist] Group: {} ClientId: {} updateConsumeOffsetToBroker {} {}",
                    this.groupName,
                    this.mQClientFactory.getClientId(),
                    mq,
                    offset.get());
            } catch (Exception e) {
                log.error("updateConsumeOffsetToBroker exception, " + mq.toString(), e);
            }
        }
    }

    /**
     * 从本地偏移量表中移除指定MQ的偏移量信息
     * @param mq 需要移除偏移量信息的MQ
     */
    public void removeOffset(MessageQueue mq) {
        if (mq != null) {
            this.offsetTable.remove(mq);
            log.info("remove unnecessary messageQueue offset. group={}, mq={}, offsetTableSize={}", this.groupName, mq,
                offsetTable.size());
        }
    }

    /**
     * 克隆指定主题的偏移量信息
     * @param topic 需要克隆偏移量的主题名称
     * @return 指定主题的克隆后的偏移量表
     */
    @Override
    public Map<MessageQueue, Long> cloneOffsetTable(String topic) {
        Map<MessageQueue, Long> cloneOffsetTable = new HashMap<>(this.offsetTable.size(), 1);

        // 遍历本地偏移量表
        for (Map.Entry<MessageQueue, AtomicLong> entry : this.offsetTable.entrySet()) {
            MessageQueue mq = entry.getKey();
            if (!UtilAll.isBlank(topic) && !topic.equals(mq.getTopic())) {
                continue;
            }
            // 将对应的MQ和偏移量信息添加到结果集合中
            cloneOffsetTable.put(mq, entry.getValue().get());
        }
        // 返回结果集合
        return cloneOffsetTable;
    }

    /**
     * 通过单向消息将指定MQ的偏移量信息更新到远程broker。
     * 一旦远程broker的master宕机，则需要更新到slave节点。
     * here need to be optimized.
     */
    private void updateConsumeOffsetToBroker(MessageQueue mq, long offset) throws RemotingException,
        MQBrokerException, InterruptedException, MQClientException {
        updateConsumeOffsetToBroker(mq, offset, true);
    }

    /**
     * 以同步的方式向远程broker更新指定mq的偏移量。
     * Once the Master is off, updated to Slave, here need to be optimized.
     */
    @Override
    public void updateConsumeOffsetToBroker(MessageQueue mq, long offset, boolean isOneway) throws RemotingException,
        MQBrokerException, InterruptedException, MQClientException {
        // 查找指定brokerName的master节点
        FindBrokerResult findBrokerResult = this.mQClientFactory.findBrokerAddressInSubscribe(this.mQClientFactory.getBrokerNameFromMessageQueue(mq), MixAll.MASTER_ID, false);

        // 如果没有找到，则从nameserver更新本地缓存的主题路由表
        if (null == findBrokerResult) {
            this.mQClientFactory.updateTopicRouteInfoFromNameServer(mq.getTopic());
            // 重新查找
            findBrokerResult = this.mQClientFactory.findBrokerAddressInSubscribe(this.mQClientFactory.getBrokerNameFromMessageQueue(mq), MixAll.MASTER_ID, false);
        }

        // 如果找到了
        if (findBrokerResult != null) {
            // 创建更新消费者偏移量的请求报文头
            UpdateConsumerOffsetRequestHeader requestHeader = new UpdateConsumerOffsetRequestHeader();
            // 设置主题名称
            requestHeader.setTopic(mq.getTopic());
            // 设置消费组名称
            requestHeader.setConsumerGroup(this.groupName);
            // 设置QueueId
            requestHeader.setQueueId(mq.getQueueId());
            // 设置偏移量
            requestHeader.setCommitOffset(offset);
            // 设置broker名称
            requestHeader.setBname(mq.getBrokerName());

            if (isOneway) {
                // 单向请求更新
                this.mQClientFactory.getMQClientAPIImpl().updateConsumerOffsetOneway(
                    findBrokerResult.getBrokerAddr(), requestHeader, 1000 * 5);
            } else {
                // 非单向请求
                this.mQClientFactory.getMQClientAPIImpl().updateConsumerOffset(
                    findBrokerResult.getBrokerAddr(), requestHeader, 1000 * 5);
            }
        } else {
            throw new MQClientException("The broker[" + mq.getBrokerName() + "] not exist", null);
        }
    }

    /**
     * 从远程broker获取指定MQ的偏移量
     * @param mq 指定的MQ
     * @return 偏移量
     * @throws RemotingException
     * @throws MQBrokerException
     * @throws InterruptedException
     * @throws MQClientException
     */
    private long fetchConsumeOffsetFromBroker(MessageQueue mq) throws RemotingException, MQBrokerException,
        InterruptedException, MQClientException {

        // 获取指定brokerName的master地址
        FindBrokerResult findBrokerResult = this.mQClientFactory.findBrokerAddressInSubscribe(this.mQClientFactory.getBrokerNameFromMessageQueue(mq), MixAll.MASTER_ID, true);

        if (null == findBrokerResult) {
            // 如果没有找到，则首先从nameserver更新指定主题的路由信息
            this.mQClientFactory.updateTopicRouteInfoFromNameServer(mq.getTopic());
            // 重新查找
            findBrokerResult = this.mQClientFactory.findBrokerAddressInSubscribe(this.mQClientFactory.getBrokerNameFromMessageQueue(mq), MixAll.MASTER_ID, false);
        }

        // 如果找到了
        if (findBrokerResult != null) {
            // 创建查询消费者偏移量的请求报文头
            QueryConsumerOffsetRequestHeader requestHeader = new QueryConsumerOffsetRequestHeader();
            // 设置主题名称
            requestHeader.setTopic(mq.getTopic());
            // 设置消费组名称
            requestHeader.setConsumerGroup(this.groupName);
            // 设置QueueId
            requestHeader.setQueueId(mq.getQueueId());
            // 设置broker名称
            requestHeader.setBname(mq.getBrokerName());

            // 发起请求，返回获取的偏移量
            return this.mQClientFactory.getMQClientAPIImpl().queryConsumerOffset(
                findBrokerResult.getBrokerAddr(), requestHeader, 1000 * 5);
        } else {
            throw new MQClientException("The broker[" + mq.getBrokerName() + "] not exist", null);
        }
    }
}
